# coding : UTF-8
# 作者：BingBO   时间：2022年10月17日

import cv2 as cv
import numpy as np
import math
from matplotlib import pyplot as plt


# 灰度重心法提取激光条纹重心函数
def COG(img):
    print("input img.shape: ", img.shape)
    height, width, channels = img.shape
    b, g, r = cv.split(img)

    ret1, th1 = cv.threshold(r, 190, 255, cv.THRESH_BINARY)  # ret1分割阈值
    th1 = cv.medianBlur(th1, 3)  # 中值滤波
    # cv.imshow('threshold img', th1)

    cp_row, cp_col, Zero_col = [], [], []
    for j in range(width):  # 遍历每一列
        x0, y0 = 0, 0
        for i in range(height):  # 遍历每一行
            x0 += th1[i, j]
            y0 += th1[i, j] * (i + 1)
        if x0 != 0:
            x0_float, y0_float = float(x0), float(y0)  # 整数相除，小数直接抛弃
            point_row = round(y0_float / x0_float) - 1  #
            # 列表存储条纹中心坐标
            cp_row += [point_row]
            cp_col += [j]  # 中心点列索引列表(包含空白区域列坐标的突变)
        else:
            Zero_col += [j]  # 缺陷点的列索引(包括了交叉光条右侧 的缺少光条空白区域)

    print("CP_row:", cp_row)
    print("CP_col:", cp_col)
    print("Zero_col:", Zero_col)

    # 创建新图像，显示中心点
    cpImg = np.zeros_like(img)
    for i in range(0, len(cp_row)):
        img[cp_row[i], cp_col[i]] = [255, 0, 0]  # 蓝色点显示
        cpImg[cp_row[i], cp_col[i]] = [255, 255, 255]
    cv.imshow("centralPointImg:", cpImg)
    # cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\centralPoint.jpg", centralPointImg)

    gapNums = []
    count = 0
    T = 20
    for i in range(0, len(cp_col) - 1):  # 中心点的列总数len(CP_col)，索引是0~len(CP_col)-1，下语句后面减去前面，需再减1
        if (cp_col[i + 1] - cp_col[i]) > T:  # 中心点列坐标的突变超过阈值50，认为是边缘大断点处，排除偶尔光条的间断
            # 在原图像上标记左右两断点
            img[cp_row[i], cp_col[i]] = [0, 255, 0]
            img[cp_row[i + 1], cp_col[i + 1]] = [0, 255, 0]
            # (左断点(含中心点)列坐标 , 右断点(含中心点)列坐标 , 左右断点间距)
            gapNums.append((cp_col[i], cp_col[i + 1], (cp_col[i + 1] - cp_col[i] - 1)))
            index = i  # 断裂处的列索引
            count += 1
    # 在原输入图像上显示断点
    cv.imshow('BreakPoint_V', img)
    # cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\BreakPoint_V.png", img)  # 保存图像
    print("共发现间距大于%d像素的(竖直方向)间断数：" % T, count)
    print("gapNums Info:", gapNums)

    # *************取断点左边点集****************
    cpLeft_row = cp_row[:index + 1]
    cpLeft_col = cp_col[:index + 1]

    for i, gapNum in enumerate(gapNums, start=1):
        print("列索引%d(左断点列坐标 , 右断点列坐标，断点间距): " % i, gapNum)

    return cpImg, cpLeft_row, cpLeft_col


# 对断点左边 中心点图像，进行 霍夫线变换
# 输入的是 中心点二值图像
def HoughLines(cpImg, cpLeft_row, cpLeft_col):
    print("\n**********************HoughLines**********************")
    gray = cv.cvtColor(cpImg, cv.COLOR_BGR2GRAY)  # 先转换为灰度图
    _, binary = cv.threshold(gray, 100, 255, cv.THRESH_BINARY)  # 必须是单通道灰度图像

    lines = cv.HoughLines(binary, 1, np.pi / 180, 150)  # 设置阈值为100
    # 返回一个：math:(rho，theta)值的数组
    print("检测到直线数量：", lines.shape[0])  # (1,1,2)
    print("lines:", lines)

    for line in lines:
        rho, theta = line[0]  # rho 是距图像左上角 (0,0) 的距离，theta 为弧度表示的线旋转角度。
        print("主线角度：", theta * 180 / 3.14)

        cos, sin = np.cos(theta), np.sin(theta)
        b = rho / sin
        k = np.tan(theta - 3.1416 / 2)
        print("k , b:", k, b)

        x1, x2 = cpLeft_col[0], cpLeft_col[len(cpLeft_col) - 1]
        y1, y2 = int(b + k * x1), int(b + k * x2)

        cv.line(cpImg, (x1, y1), (x2, y2), (0, 255, 0), 1)
        cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\HoughLines.png",
                   cpImg)  # 保存图像

    # 计算各个中心点到主线的距离，找出最远特征点的坐标
    dMax, index = 0, 0
    for i in range(0, len(cpLeft_row)):
        d = math.fabs(k * cpLeft_col[i] - cpLeft_row[i] + b) / math.sqrt(1 + k * k)  # 主线方程：kx-y+b=0，(A=k B=-1 C=b)
        if d >= dMax:
            dMax = d
            coor = (cpLeft_row[i], cpLeft_col[i])
            index = i
    print("距离主线最远中心点坐标：", coor)
    cpImg[coor[0], coor[1]] = [0, 0, 255]  # 标记最远点
    cv.imshow("FurthestPoint", cpImg)
    # cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\furthestPoint.png", img)  # 保存图像

    # ***********分割最远点左右点集*************
    fpLeft_row, fpLeft_col = cpLeft_row[:index + 1], cpLeft_col[:index + 1]
    fpRight_row, fpRight_col = cpLeft_row[index + 1:], cpLeft_col[index + 1:]

    # 创建单通道左中心点图像
    fpLeftImg = np.zeros_like(cpImg)
    for i in range(0, len(fpLeft_row)):
        fpLeftImg[fpLeft_row[i], fpLeft_col[i]] = [255, 255, 255]
    cv.imshow("fpLeftImg", fpLeftImg)
    # cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\fpLeftImg.jpg", fpLeftImg)

    # 创建单通道右中心点图像
    fpRightImg = np.zeros_like(cpImg)
    for i in range(0, len(fpRight_row)):
        fpRightImg[fpRight_row[i], fpRight_col[i]] = [255, 255, 255]
    cv.imshow("fpRightImg:", fpRightImg)
    # cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\fpLeftImg.jpg", fpLeftImg)

    return fpLeftImg, fpRightImg, fpLeft_row, fpLeft_col, fpRight_row, fpRight_col


def HoughLines_Left(fpLeftImg, fpLeft_row, fpLeft_col):
    # ***************1. 对左半段单通道图像hough直线拟合******************
    print("\n********************HoughLines_Left************************")
    gray = cv.cvtColor(fpLeftImg, cv.COLOR_BGR2GRAY)  # 先转换为灰度图
    _, binary = cv.threshold(gray, 100, 255, cv.THRESH_BINARY)  # 必须是单通道灰度图像

    fpLeft_mainline = cv.HoughLines(binary, 1, np.pi / 180, 100)  # 设置阈值为100
    print(fpLeft_mainline.shape)  # (1,1,2)
    print("linesLeft(rho，theta):", fpLeft_mainline)

    for lineLeft in fpLeft_mainline:
        rho, theta = lineLeft[0]
        print("左半部主线角度：", theta * 180 / 3.14)

        cos, sin = np.cos(theta), np.sin(theta)
        b = rho / sin
        k = np.tan(theta - 3.1416 / 2)
        print("左半部主线k , b:", k, b)

        x1, x2 = fpLeft_col[0], fpLeft_col[len(fpLeft_col) - 1]
        y1, y2 = int(b + k * x1), int(b + k * x2)

        cv.line(fpLeftImg, (x1, y1), (x2, y2), (0, 255, 0), 1)
    cv.imshow("mainline of fpLeftImg", fpLeftImg)
    # cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\img_new_Left.png", img_new_Left)

    # 分割出 距离主线>阈值T的  点集
    T, index, count = 2, 0, 0
    grooveLeftImg = np.zeros_like(fpLeftImg)
    for i in range(0, len(fpLeft_row)):
        d = math.fabs(k * fpLeft_col[i] - fpLeft_row[i] + b) / math.sqrt(1 + k * k)
        # print("d:", d)
        if d >= T:
            grooveLeftImg[fpLeft_row[i], fpLeft_col[i]] = [255, 255, 255]
            count += 1
    print("共找到大于阈值T的点数：", count)
    cv.imshow("grooveLeftImg", grooveLeftImg)
    print("grooveLeftImg.shape：", grooveLeftImg.shape)

    # ***************2.对左侧坡口上的图片点再进行Hough直线检测，返回该直线参数******************
    gray = cv.cvtColor(grooveLeftImg, cv.COLOR_BGR2GRAY)  # 先转换为灰度图
    _, binary = cv.threshold(gray, 100, 255, cv.THRESH_BINARY)  # 必须是单通道灰度图像

    linesOnGrooveLeft = cv.HoughLines(binary, 1, np.pi / 180, 30)  # 设置阈值为30
    print(linesOnGrooveLeft.shape)  # (1,1,2)
    print("linesOnGrooveLeft(rho，theta):", linesOnGrooveLeft)  # 获取主线参数1

    for lineLeft in linesOnGrooveLeft:
        rho, theta = lineLeft[0]
        print("左半部坡口主线角度：", theta * 180 / 3.14)

        a, b = np.cos(theta), np.sin(theta)
        x0, y0 = a * rho, b * rho  # 原点到直线的 垂点坐标
        lenL, lenR = 500, 600
        x1, y1 = int(x0 + lenL * (-b)), int(y0 + lenL * (a))
        x2, y2 = int(x0 - lenR * (-b)), int(y0 - lenR * (a))
        # 获取主线参数2
        k = (y2 - y1) / (x2 - x1)
        b = y1 - k * x1
        print("左半部坡口斜率k: %f" % k, "截距b：%f" % b)
        cv.line(fpLeftImg, (x1, y1), (x2, y2), (0, 255, 0), 1)
    # cv.imshow("mainline of pointOnGrooveLeft", pointOnGrooveLeft)  # 显示不了主线，则取到主线参数也可以
    return linesOnGrooveLeft, k, b


def HoughLines_Right(fpRightImg, fpRight_row, fpRight_col):
    # ***************1.对右半段单通道图像hough直线拟合******************
    print("\n*********************HoughLines_Right***********************")
    gray = cv.cvtColor(fpRightImg, cv.COLOR_BGR2GRAY)  # 先转换为灰度图
    _, binary = cv.threshold(gray, 100, 255, cv.THRESH_BINARY)  # 必须是单通道灰度图像

    linesRight = cv.HoughLines(binary, 1, np.pi / 180, 100)  # 设置阈值为100
    print(linesRight.shape)  # (1,1,2)
    print("linesRight(rho，theta):", linesRight)

    for lineRight in linesRight:
        rho, theta = lineRight[0]
        print("右半部主线角度：", theta * 180 / 3.14)

        cos, sin = np.cos(theta), np.sin(theta)
        b = rho / sin
        k = np.tan(theta - 3.1416 / 2)
        print("右半部主线k , b:", k, b)

        x1, x2 = fpRight_col[0], fpRight_col[len(fpRight_col) - 1]
        y1, y2 = int(b + k * x1), int(b + k * x2)

        cv.line(fpRightImg, (x1, y1), (x2, y2), (0, 255, 0), 1)
    cv.imshow("mainline of fpRightImg", fpRightImg)
    # cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\img_new_Right.png", img_new_Right)

    # 分割出 距离主线>阈值T的  点集
    T, index, count = 2, 0, 0
    grooveRightImg = np.zeros_like(fpRightImg)
    for i in range(0, len(fpRight_row)):
        d = math.fabs(k * fpRight_col[i] - fpRight_row[i] + b) / math.sqrt(1 + k * k)
        # print("d:", d)
        if d >= T:
            grooveRightImg[fpRight_row[i], fpRight_col[i]] = [255, 255, 255]
            count += 1
            # coor = (furthestPoint_row_Left[i], furthestPoint_col_Left[i])
            # index = i
    print("右坡口共找到大于阈值T的点数：", count)
    cv.imshow("grooveRightImg", grooveRightImg)
    print("pointOnGrooveRight：", grooveRightImg.shape)

    # ***************2.对右侧坡口上的图片点再进行Hough直线检测，返回该直线参数******************
    gray = cv.cvtColor(grooveRightImg, cv.COLOR_BGR2GRAY)  # 先转换为灰度图
    _, binary = cv.threshold(gray, 100, 255, cv.THRESH_BINARY)  # 必须是单通道灰度图像

    linesOnGrooveRight = cv.HoughLines(binary, 1, np.pi / 180, 14)  # 设置阈值为14
    print(linesOnGrooveRight.shape)  # (1,1,2)
    print("linesOnGrooveLeft(rho，theta):", linesOnGrooveRight)  # 获取主线参数1

    for lineRight in linesOnGrooveRight:
        rho, theta = lineRight[0]
        print("右半部坡口主线角度：", theta * 180 / 3.14)

        a, b = np.cos(theta), np.sin(theta)
        x0, y0 = a * rho, b * rho  # 原点到直线的 垂点坐标
        lenL, lenR = 500, 600
        x1, y1 = int(x0 + lenL * (-b)), int(y0 + lenL * (a))
        x2, y2 = int(x0 - lenR * (-b)), int(y0 - lenR * (a))
        # 获取主线参数2
        k = (y2 - y1) / (x2 - x1)
        b = y1 - k * x1
        print("右边半部坡口斜率k: %f" % k, "截距b：%f" % b)
        cv.line(fpRightImg, (x1, y1), (x2, y2), (0, 255, 0), 1)
    cv.imshow("mainline of groovePointRight", grooveRightImg)  # 显示不了主线，则取到主线参数也可以
    return linesOnGrooveRight, k, b


def featurePoint(img, k1, b1, k2, b2):
    x = (b2 - b1) / (k1 - k2)  # 列坐标
    y = k1 * x + b1  # 行坐标
    # a, b = int(x), int(y)  # 直接省略小数部分
    row, col = round(y), round(x)
    img[row, col] = [0, 255, 0]

    cv.imshow("FeaturePoint", img)
    cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\FeaturePoint_Src.png", img)
    return row, col


src = cv.imread(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\V\Image_20220909101553637.jpg")
src = cv.resize(src, None, fx=0.4, fy=0.4, interpolation=cv.INTER_CUBIC)

cpImg, cpLeft_row, cpLeft_col = COG(src)

cpLeftImg, cpRightImg, fpLeft_row, fpLeft_col, fpRight_row, fpRight_col = HoughLines(cpImg, cpLeft_row, cpLeft_col)

linesOnGrooveLeft, k1, b1 = HoughLines_Left(cpLeftImg, fpLeft_row, fpLeft_col)
linesOnGrooveRight, k2, b2 = HoughLines_Right(cpRightImg, fpRight_row, fpRight_col)
featurePoint(src, k1, b1, k2, b2)

cv.waitKey(0)
cv.destroyAllWindows()
